package ch.ethz.fcl.metrobuzz.render.od.flow_view.sankey;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import ch.ethz.fcl.metrobuzz.data.od.ODMatrixElement;
import ch.ethz.fcl.metrobuzz.data.od.SelectLinkQuery;
import ch.ethz.fcl.metrobuzz.data.scene.basemap.Node;

public class ODSankeyDiagramData {
	private SelectLinkQuery slq;

	private List<SankeyNode> allOriginSNList;
	private List<SankeyNode> allDestinSNList;

	private List<SankeyNodeGroup> initOGroupList;
	private List<SankeyNodeGroup> initDGroupList;

	private List<SankeyNodeGroup> curOGroupList;
	private List<SankeyNodeGroup> curDGroupList;

	private Map<Node, SankeyNode> startNodeSankeyMap;
	private Map<Node, SankeyNode> endNodeSankeyMap;

	private Map<SankeyNode, SankeyNodeGroup> curStartNodeGroupMap;
	private Map<SankeyNode, SankeyNodeGroup> initStartNodeGroupMap;

	private Map<SankeyNode, SankeyNodeGroup> curEndNodeGroupMap;
	private Map<SankeyNode, SankeyNodeGroup> initEndNodeGroupMap;

	// private Map<ODPair, Integer> sankeyPairMap;
	private List<ODPair> odPairList;

	public static boolean time_changed = false;
	public static boolean reset = false;

	public ODSankeyDiagramData(SelectLinkQuery slq) {
		this.slq = slq;
	}

	public void update() {
		if (startNodeSankeyMap == null || !time_changed || reset) {
			startNodeSankeyMap = new HashMap<Node, SankeyNode>();
			endNodeSankeyMap = new HashMap<Node, SankeyNode>();

			initStartNodeGroupMap = new HashMap<SankeyNode, SankeyNodeGroup>();
			initEndNodeGroupMap = new HashMap<SankeyNode, SankeyNodeGroup>();

			reset = false;
		} else if (time_changed) {
			resetVolume(new ArrayList<SankeyNodeGroup>(
					initStartNodeGroupMap.values()));
			resetVolume(new ArrayList<SankeyNodeGroup>(
					initEndNodeGroupMap.values()));

			time_changed = false;
		}

		curStartNodeGroupMap = new HashMap<SankeyNode, SankeyNodeGroup>();
		curEndNodeGroupMap = new HashMap<SankeyNode, SankeyNodeGroup>();

		allOriginSNList = generateSankeyNodes(slq.getODMatrix()
				.getAllOriginMatrix(), true);
		allDestinSNList = generateSankeyNodes(slq.getODMatrix()
				.getAllDestinMatrix(), false);

		initOGroupList = groupSankeyNodesByTransitLine(allOriginSNList);
		initDGroupList = groupSankeyNodesByTransitLine(allDestinSNList);

		curOGroupList = expandGroups(initOGroupList);
		curDGroupList = expandGroups(initDGroupList);

		for (SankeyNodeGroup sng : curOGroupList)
			sng.resetNodeMap();

		for (SankeyNodeGroup sng : curDGroupList)
			sng.resetNodeMap();

		generateSankeyPair();
		generateSankeyPairMap();
	}

	private void resetVolume(List<SankeyNodeGroup> nodeGroupList) {
		for (SankeyNodeGroup sng : nodeGroupList) {
			for (SankeyNode sn : sng.getNodeList())
				sn.resetVolume();
			sng.resetVolume();
			if (sng.getNodeGroupList() != null)
				resetVolume(sng.getNodeGroupList());
		}
	}

	public void expandGroups() {
		curOGroupList = expandGroups(initOGroupList);
		curDGroupList = expandGroups(initDGroupList);

		for (SankeyNodeGroup sng : curOGroupList)
			sng.resetNodeMap();

		for (SankeyNodeGroup sng : curDGroupList)
			sng.resetNodeMap();

		generateSankeyPair();
		generateSankeyPairMap();

		for (SankeyNodeGroup sng : curOGroupList)
			sng.calODPairBound();

		for (SankeyNodeGroup sng : curDGroupList)
			sng.calODPairBound();
	}

	private List<SankeyNode> generateSankeyNodes(Map<Node, Float>[] allODs,
			boolean origin) {
		List<SankeyNode> sankeyNodes = new ArrayList<SankeyNode>();

		for (int k = 0; k < allODs.length; k++) {
			Map<Node, Float> curOD = allODs[k];
			if (curOD == null || curOD.size() == 0)
				continue;

			List<Node> allNodes = new ArrayList<Node>(curOD.keySet());
			for (int i = 0; i < allNodes.size(); i++) {
				Node n = allNodes.get(i);
				Float volume = curOD.get(n);

				SankeyNode sankeyNode = null;
				if (origin)
					sankeyNode = startNodeSankeyMap.get(n);
				else
					sankeyNode = endNodeSankeyMap.get(n);

				if (sankeyNode == null) {
					sankeyNode = new SankeyNode(this, n.getId(), origin);
				} else
					sankeyNode.resetVolume(k);

				sankeyNode.addVolume(k, volume);

				if (!sankeyNodes.contains(sankeyNode))
					sankeyNodes.add(sankeyNode);

				if (origin && startNodeSankeyMap.get(n) == null)
					startNodeSankeyMap.put(n, sankeyNode);
				else if (!origin && endNodeSankeyMap.get(n) == null)
					endNodeSankeyMap.put(n, sankeyNode);
			}
		}
		return sankeyNodes;
	}

	private List<SankeyNodeGroup> expandGroups(List<SankeyNodeGroup> sngList) {
		List<SankeyNodeGroup> curList = new ArrayList<SankeyNodeGroup>();
		for (SankeyNodeGroup sng : sngList) {
			List<SankeyNodeGroup> subSNGList = expandSankeyNodeGroup(sng);
			for (SankeyNodeGroup subSNG : subSNGList) {
				curList.add(subSNG);

				for (SankeyNode sn : subSNG.getNodeList()) {
					if (sn.origin)
						curStartNodeGroupMap.put(sn, subSNG);
					else
						curEndNodeGroupMap.put(sn, subSNG);
				}
			}
		}

		return curList;
	}

	private List<SankeyNodeGroup> expandSankeyNodeGroup(SankeyNodeGroup sng) {
		List<SankeyNodeGroup> sngList = new ArrayList<SankeyNodeGroup>();
		if (sng.hide)
			return sngList;

		if (sng.expanded()) {
			for (SankeyNodeGroup subSNG : sng.getNodeGroupList()) {
				List<SankeyNodeGroup> subSNGList = expandSankeyNodeGroup(subSNG);
				for (SankeyNodeGroup ss : subSNGList)
					sngList.add(ss);
			}
		} else {
			sngList.add(sng);
		}
		return sngList;
	}

	private void generateSankeyPair() {
		Map<ODMatrixElement, Float>[] allODPairs = slq.getODMatrix()
				.getAllODMatrix();
		for (int time = 0; time < allODPairs.length; time++) {
			Map<ODMatrixElement, Float> odPairs = allODPairs[time];
			if (odPairs == null || odPairs.size() == 0)
				continue;

			for (ODMatrixElement od : odPairs.keySet()) {
				Node origin = od.getOrigin();
				Node destin = od.getDestination();

				float edgeODProb = slq
						.getData()
						.getEdgeODPoss()
						.getEdgeODPoss(slq.getStart(), slq.getEnd(), origin,
								destin);

				Float volume = odPairs.get(od);

				SankeyNode start = startNodeSankeyMap.get(od.getOrigin());
				SankeyNode end = endNodeSankeyMap.get(od.getDestination());

				Float startToEnd = start.getNodePairVolume(time, end);
				if (startToEnd == null) {
					start.putNodePair(time, end, volume);
				} else {
					start.putNodePair(time, end, volume + startToEnd);
				}

				Float endFromStart = end.getNodePairVolume(time, start);
				if (endFromStart == null) {
					end.putNodePair(time, start, volume);
				} else {
					end.putNodePair(time, start, volume + endFromStart);
				}

				SankeyNodeGroup startGroup = curStartNodeGroupMap.get(start);
				SankeyNodeGroup endGroup = curEndNodeGroupMap.get(end);

				Float startGroupToEndGroup = startGroup.getNodePairVolume(
						time, endGroup);
				if (startGroupToEndGroup == null) {
					startGroup.putNodePair(time, endGroup, volume);
				} else {
					startGroup.putNodePair(time, endGroup, volume
							+ startGroupToEndGroup);
				}

				Float nodeProb = startGroup.getNodeProb(endGroup);
				if (nodeProb == null) {
					startGroup.putNodeProb(endGroup, edgeODProb);

					// if (startGroup.getId().equals("EW26")
					// && endGroup.getId().equals("EW15"))
					// System.out.println("edge prob " + edgeODProb
					// + " pre null");
				} else {
					float size = startGroup.getNodePairVolume(time, endGroup);
					float newProb = (volume * edgeODProb + nodeProb * size)
							/ (volume + size);
					startGroup.putNodeProb(endGroup, newProb);

					// if (startGroup.getId().equals("EW26")
					// && endGroup.getId().equals("EW15"))
					// System.out.println("edge prob " + edgeODProb
					// + " pre " + nodeProb);
				}

				Float endGroupFromStartGroup = endGroup.getNodePairVolume(
						time, startGroup);
				if (endGroupFromStartGroup == null)
					endGroup.putNodePair(time, startGroup, volume);
				else
					endGroup.putNodePair(time, startGroup, volume
							+ endGroupFromStartGroup);
			}
		}
	}

	protected List<SankeyNodeGroup> groupSankeyNodesByItself(
			List<SankeyNode> sankeyNodeList) {
		List<SankeyNodeGroup> nodeGroupList = new ArrayList<SankeyNodeGroup>();

		for (SankeyNode sn : sankeyNodeList) {
			SankeyNodeGroup snGroup = new SankeyNodeGroup(this, sn.getId(),
					sn.origin);

			int[] allVolumes = sn.getAllVolume();
			for (int i = 0; i < allVolumes.length; i++)
				snGroup.addVolume(i, allVolumes[i]);

			snGroup.addSankeyNode(sn);

			if (sn.origin)
				curStartNodeGroupMap.put(sn, snGroup);
			else
				curEndNodeGroupMap.put(sn, snGroup);

			nodeGroupList.add(snGroup);
		}

		return nodeGroupList;
	}

	protected List<SankeyNodeGroup> groupSankeyNodesByTransitLine(
			List<SankeyNode> nodeList) {
		List<SankeyNodeGroup> nodeGroupList = new ArrayList<SankeyNodeGroup>();

		for (int i = 0; i < nodeList.size(); i++) {
			SankeyNode sn = nodeList.get(i);

			SankeyNodeGroup snGroup = null;
			if (sn.origin)
				snGroup = initStartNodeGroupMap.get(sn);
			else
				snGroup = initEndNodeGroupMap.get(sn);

			if (snGroup == null) {
				String lineName = sn.id.substring(0, 2);
				if (lineName.equals("CG"))
					lineName = "EW";

				else if (lineName.equals("BP"))
					lineName = "NS";

				for (int j = 0; j < nodeGroupList.size(); j++) {
					if (nodeGroupList.get(j).id.equals(lineName)) {
						snGroup = nodeGroupList.get(j);
						break;
					}
				}

				if (snGroup == null) {
					snGroup = new SankeyNodeGroup(this, lineName, sn.origin);
					nodeGroupList.add(snGroup);
				}

				if (sn.origin) {
					initStartNodeGroupMap.put(sn, snGroup);
					curStartNodeGroupMap.put(sn, snGroup);
				} else {
					initEndNodeGroupMap.put(sn, snGroup);
					curEndNodeGroupMap.put(sn, snGroup);
				}
			} else {
				if (!nodeGroupList.contains(snGroup))
					nodeGroupList.add(snGroup);

				if (sn.origin)
					curStartNodeGroupMap.put(sn, snGroup);
				else
					curEndNodeGroupMap.put(sn, snGroup);
			}

			int[] allVolumes = sn.getAllVolume();

			for (int k = 0; k < allVolumes.length; k++)
				snGroup.addVolume(k, allVolumes[k]);

			snGroup.addSankeyNode(sn);
		}

		SankeyNode.sortByVolume(nodeGroupList);

		for (SankeyNodeGroup sng : nodeGroupList) {
			if (sng.bound != null)
				sng.generateSubGroup();
		}

		return nodeGroupList;
	}

	protected List<SankeyNodeGroup> groupSankeyNodesByVolume(
			List<SankeyNode> nodeList) {
		List<SankeyNodeGroup> nodeGroupList = new ArrayList<SankeyNodeGroup>();

		float totalNumber = slq.getPassTripSize();
		
		@SuppressWarnings("unused")
		float aggRatio = 0;

		for (int i = 0; i < nodeList.size(); i++) {
			SankeyNode sn = nodeList.get(i);

			if (nodeGroupList.size() >= 4) {
				SankeyNodeGroup sankeyNodeGroup = new SankeyNodeGroup(this,
						"Others", sn.origin);

				for (int j = i; j < nodeList.size(); j++) {
					int[] allVolumes = nodeList.get(j).getAllVolume();
					for (int k = 0; k < allVolumes.length; k++)
						sankeyNodeGroup.addVolume(k, allVolumes[k]);

					sankeyNodeGroup.addSankeyNode(nodeList.get(j));

					if (nodeList.get(j).origin)
						curStartNodeGroupMap.put(nodeList.get(j),
								sankeyNodeGroup);
					else
						curEndNodeGroupMap
								.put(nodeList.get(j), sankeyNodeGroup);
				}

				nodeGroupList.add(sankeyNodeGroup);
				break;
			} else {
				SankeyNodeGroup sankeyNodeGroup = new SankeyNodeGroup(this,
						sn.getId(), sn.origin);

				int[] allVolumes = sn.getAllVolume();
				for (int k = 0; k < allVolumes.length; k++)
					sankeyNodeGroup.addVolume(k, allVolumes[k]);

				sankeyNodeGroup.addSankeyNode(nodeList.get(i));

				if (sn.origin)
					curStartNodeGroupMap.put(sn, sankeyNodeGroup);
				else
					curEndNodeGroupMap.put(sn, sankeyNodeGroup);

				nodeGroupList.add(sankeyNodeGroup);
			}

			aggRatio += sn.getVolume() * 1.0f / totalNumber;
		}

		return nodeGroupList;
	}

	private void generateSankeyPairMap() {
		odPairList = new ArrayList<ODPair>();
		for (SankeyNodeGroup origin : curOGroupList) {
			Map<SankeyNode, Float> destinSNMap = origin.getCurODPairMap();
			for (SankeyNode destin : destinSNMap.keySet()) {
				ODPair pair = new ODPair(origin, destin);
				pair.setVolume(destinSNMap.get(destin));
				Float prob = origin.getNodeProb(destin);
				if (prob == null)
					pair.setProbability(0);
				else
					pair.setProbability(prob);
				odPairList.add(pair);
			}
		}
	}

	public List<SankeyNodeGroup> getCurOrginSNList() {
		return curOGroupList;
	}

	public List<SankeyNodeGroup> getCurDestinSNList() {
		return curDGroupList;
	}

	public List<SankeyNodeGroup> getInitOrginSNList() {
		return initOGroupList;
	}

	public List<SankeyNodeGroup> getInitDestinSNList() {
		return initDGroupList;
	}

	public List<ODPair> getODPairs() {
		return odPairList;
	}

	public SelectLinkQuery getSLQ() {
		return slq;
	}
}
